/* 
 * Enhanced Correlation Coefficient (ECC) Image Alignment Algorithm 
 * http://xanthippi.ceid.upatras.gr/people/evangelidis/ecc
 *
 * Authors: Georgios Evangelidis
 * e-mail: evagelid@ceid.upatras.gr
 *
 * Copyright (2010): G.D.Evangelidis
 *
 * For details, see the paper:
 * G.D. Evangelidis, E.Z. Psarakis, "Parametric Image Alignment using 
 * Enhanced Correlation Coefficient Maximization, IEEE Transaction on
 * Pattern Analysis & Machine Intelligence, vol. 30, no. 10, 2008
 */

#include <cassert>
#include <opencv/cv.h>
#include "eccBlock.h"

static void update_warping_matrix (CvMat * map_matrix, const CvMat * update);

static void image_jacobian_homo(CvMat * src1, CvMat * src2, const CvMat * src3, const CvMat * src4, const CvMat * src5, CvMat * dst);
static void image_jacobian_affine(const CvMat * src1, const CvMat * src2, const CvMat * src3, const CvMat * src4, CvMat * dst);
static void project_onto_jacobian(const CvMat * src1, const CvMat * src2, CvMat * dst);


using namespace std;


static void project_onto_jacobian(const CvMat * src1, const CvMat * src2, CvMat * dst)
{

	assert(src1);
	assert(src2);
	assert(dst);
	assert(src1->height == src2->height);
	assert((src1->width % src2->width) == 0);
	int w;
	CvMat mat;

	if (dst->width == 1){
		w  = src2->width;
		for (int i=0; i<dst->height; i++){
			cvGetCols(src1, &mat, i*w, (i+1)*w);
			CV_MAT_ELEM(*dst, float, i, 0) = cvDotProduct(src2, &mat);
		}
	}

	else {
		assert(dst->width == dst->height); //dst is symmetric
		w = src2->width/dst->width;
		CvMat mat2;
		for (int i=0; i<dst->height; i++){
			cvGetCols(src1, &mat, i*w, (i+1)*w);
			CV_MAT_ELEM(*dst, float, i, i) = cvDotProduct(&mat, &mat); //diagonal elements
			for (int j=i+1; j<dst->width; j++){ //j starts from i+1
				cvGetCols(src2, &mat2, j*w, (j+1)*w);
				const double cross = cvDotProduct(&mat, &mat2); 
				CV_MAT_ELEM(*dst, float, i, j) = cross; //above diagonal
					CV_MAT_ELEM(*dst, float, j, i) = cross; //below diagonal
			}
		}
	}
}


static void image_jacobian_homo(CvMat * src1, CvMat * src2, 
								const CvMat * src3, const CvMat * src4,
								const CvMat * src5, CvMat * dst)
{
	assert(src1);
	assert(src2);
	assert(src3);
	assert(src4);
	assert(src5);
	assert(dst);
	
	const float h0 = CV_MAT_ELEM(*src5, float, 0, 0);  
	const float h1 = CV_MAT_ELEM(*src5, float, 1, 0);  
	const float h2 = CV_MAT_ELEM(*src5, float, 2, 0);  
	const float h3 = CV_MAT_ELEM(*src5, float, 0, 1);  
	const float h4 = CV_MAT_ELEM(*src5, float, 1, 1);  
	const float h5 = CV_MAT_ELEM(*src5, float, 2, 1);  
	const float h6 = CV_MAT_ELEM(*src5, float, 0, 2);  
	const float h7 = CV_MAT_ELEM(*src5, float, 1, 2);  

	CvMat mat;
	const int w = src1->width;

	CvMat * hatX = cvCreateMat(src1->height, src1->width, CV_32F);
	CvMat * hatY = cvCloneMat(hatX);
	CvMat * den = cvCloneMat(hatX);
	CvMat * temp = cvCloneMat(hatX);


	//create denominator for all points as a block 
	cvAddWeighted(src3, h2, src4, h5, 1.0, den);
	
	//create projected points similarly (you need the negative)
	cvAddWeighted(src3, -h0, src4, -h3, -h6, hatX);
	cvDiv(hatX, den, hatX);
	cvAddWeighted(src3, -h1, src4, -h4, -h7, hatY);
	cvDiv(hatY, den, hatY);

	//instead of dividing each block with den, 
	//just pre-devide the block of gradients
	cvDiv(src1, den, src1);
	cvDiv(src2, den, src2);
	cvReleaseMat(&den);

	//compute Jacobian blocks
	cvGetCols(dst, &mat, 0, w);
	cvMul(src1, src3, &mat);

	cvGetCols(dst, &mat, w, 2*w);
	cvMul(src2, src3, &mat);

	cvGetCols(dst, &mat, 2*w, 3*w);
	cvMul(hatX, src1, temp);
	cvMul(hatY, src2, &mat);
	cvAdd(temp, &mat, temp);
	cvMul(temp, src3, &mat);

	cvReleaseMat(&hatX);
	cvReleaseMat(&hatY);

	cvGetCols(dst, &mat, 3*w, 4*w);
	cvMul(src1, src4, &mat);

	cvGetCols(dst, &mat, 4*w, 5*w);
	cvMul(src2, src4, &mat);

	cvGetCols(dst, &mat, 5*w, 6*w);
	cvMul(temp, src4, &mat);

	cvGetCols(dst, &mat, 6*w, 7*w);
	cvCopy(src1, &mat);

	cvGetCols(dst, &mat, 7*w, 8*w);
	cvCopy(src2, &mat);

	cvReleaseMat(&temp);	
}


static void image_jacobian_affine(const CvMat * src1, const CvMat * src2, 
								const CvMat * src3, const CvMat * src4,
								CvMat * dst)
{
	assert(src1);
	assert(src2);
	assert(src3);
	assert(src4);
	assert(dst);
	
	CvMat mat;
	const int w = src1->width;

	cvGetCols(dst, &mat, 0, w);
	cvMul(src1, src3, &mat);

	cvGetCols(dst, &mat, w, 2*w);
	cvMul(src2, src3, &mat);

	cvGetCols(dst, &mat, 2*w, 3*w);
	cvMul(src1, src4, &mat);

	cvGetCols(dst, &mat, 3*w, 4*w);
	cvMul(src2, src4, &mat);

	cvGetCols(dst, &mat, 4*w, 5*w);
	cvCopy(src1, &mat);

	cvGetCols(dst, &mat, 5*w, 6*w);
	cvCopy(src2, &mat);
	
}

#pragma mark -
#pragma mark helper methods


static void update_warping_matrix (CvMat * map_matrix, const CvMat * update)
{
  assert (map_matrix); assert (CV_MAT_TYPE (map_matrix->type) == CV_32FC1);
  assert (update);     assert (CV_MAT_TYPE (update->type)     == CV_32FC1);
  assert (map_matrix->rows == 3 && update->rows == 8 || map_matrix->rows == 2 && update->rows == 6);
  assert (update->cols == 1);
  
  const int cols = map_matrix->cols;
  const int rows = map_matrix->rows;
  
  for (ptrdiff_t y = 0; y < rows; y++)
		for (ptrdiff_t x = 0; x < cols; x++)
      if (y != 2 || x != 2) // the last element (2,2) is not represented in update (homography case)
				CV_MAT_ELEM(*(map_matrix), float, y, x) += CV_MAT_ELEM(*(update), float, rows*x+y, 0);
}

#pragma mark -
#pragma mark implementation of external interface

CvMat * cvFindTransform (const IplImage * src_image, const IplImage * dst_image, CvMat * map_matrix, const WARP_MODE warp_mode, const CvTermCriteria & termination)
{
  assert (src_image); assert (src_image->nChannels == 1); assert (src_image->width > 0 && src_image->height > 0);
  assert (dst_image); assert (dst_image->nChannels == 1); assert (dst_image->width > 0 && dst_image->height > 0);
  //assert (src_image->width <= dst_image->width && src_image->height <= dst_image->height);
  assert (map_matrix); assert (map_matrix->cols == 3); assert (map_matrix->rows == 2 || map_matrix->rows ==3);
  assert (warp_mode == WARP_MODE_AFFINE || warp_mode == WARP_MODE_HOMOGRAPHY);
  assert (termination.type & CV_TERMCRIT_ITER || termination.type & CV_TERMCRIT_EPS);
  
  const int    number_of_iterations = (termination.type & CV_TERMCRIT_ITER) ? termination.max_iter : 200;
  const double termination_eps      = (termination.type & CV_TERMCRIT_EPS)  ? termination.epsilon  :  -1;

  const int number_of_model_parameters = (warp_mode == WARP_MODE_AFFINE) ? 6 : 8;
  const int number_of_pixels           = src_image->width * src_image->height;

  //form the block with x and y coordinates that corresponds to template's grid
  CvMat* x_coord = cvCreateMat(1, src_image->width, CV_32F);
  CvMat* y_coord = cvCreateMat(src_image->height, 1, CV_32F);
  CvMat* X = cvCreateMat(src_image->height, src_image->width, CV_32F);
  CvMat* Y = cvCreateMat(src_image->height, src_image->width, CV_32F);

  int j;
  for (j=0; j<src_image->height; j++)
	  CV_MAT_ELEM(*y_coord, float, j, 0) = (float)j; 

  for (j=0; j<src_image->width; j++)
	  CV_MAT_ELEM(*x_coord, float, 0, j) = (float)j; 

  cvRepeat(x_coord, X);
  cvRepeat(y_coord, Y);
  cvReleaseMat(&x_coord);
  cvReleaseMat(&y_coord);
  
  CvMat * src_image_zeromean = cvCreateMat (src_image->height, src_image->width, CV_32F); // smoothed zero-mean template
  CvMat * dst_image_float    = cvCreateMat (dst_image->height, dst_image->width, CV_32F); // smoothed input image
  CvMat * dst_image_warped   = cvCreateMat (src_image->height, src_image->width, CV_32F); // warped zero-mean input image
  
  // gaussian smoothing and zero-mean shifting
  cvConvert (src_image, src_image_zeromean);
  cvSmooth  (src_image_zeromean, src_image_zeromean, CV_GAUSSIAN, 5, 5);
  cvSubS    (src_image_zeromean, cvAvg (src_image_zeromean), src_image_zeromean);
  
  // gaussian smoothing
  cvConvert (dst_image, dst_image_float);
  cvSmooth  (dst_image_float, dst_image_float, CV_GAUSSIAN, 5, 5);
  
  // needed matrices for gradients and warped gradients
  CvMat * dst_gradient_x        = cvCreateMat (dst_image->height, dst_image->width, CV_32FC1);
  CvMat * dst_gradient_y        = cvCreateMat (dst_image->height, dst_image->width, CV_32FC1);
  CvMat * dst_gradient_x_warped = cvCreateMat (src_image->height, src_image->width, CV_32FC1);
  CvMat * dst_gradient_y_warped = cvCreateMat (src_image->height, src_image->width, CV_32FC1);
  
  // calculate first order image derivatives
  CvMat * dx = cvCreateMat (1, 3, CV_32FC1);
  CvMat * dy = cvCreateMat (3, 1, CV_32FC1);
  CV_MAT_ELEM(* dx, float, 0, 0) = -1.0f/2 ;
  CV_MAT_ELEM(* dx, float, 0, 1) =  0.0f   ;
  CV_MAT_ELEM(* dx, float, 0, 2) =  1.0f/2 ;
  CV_MAT_ELEM(* dy, float, 0, 0) = -1.0f/2 ;
  CV_MAT_ELEM(* dy, float, 1, 0) =  0.0f   ;
  CV_MAT_ELEM(* dy, float, 2, 0) =  1.0f/2 ;
  cvFilter2D (dst_image_float, dst_gradient_x, dx);
  cvFilter2D (dst_image_float, dst_gradient_y, dy);
  cvReleaseMat (& dy);
  cvReleaseMat (& dx);
  
  // linear equation system for maximizing the enhanced correlation coefficient
  CvMat * jacobian                     = cvCreateMat (src_image->height, src_image->width*number_of_model_parameters, CV_32F); // jacobian of warped image
  CvMat * hessian                      = cvCreateMat (number_of_model_parameters, number_of_model_parameters, CV_32F); // Hessian
  CvMat * hessian_inv                  = cvCloneMat  (hessian);                                                        // inverse of Hessian
  CvMat * dst_image_projected          = cvCreateMat (number_of_model_parameters, 1, CV_32F);                          // destination image projected into jacobian
  CvMat * src_image_projected          = cvCloneMat  (dst_image_projected);                                            // source image projected into jacobian
  CvMat * dst_image_projected_hessian  = cvCloneMat  (dst_image_projected);                                            // auxiliary matrix
  CvMat * error_projected              = cvCloneMat  (dst_image_projected);                                            // error projected into jacobian
  
  CvMat * delta_p_flat            = cvCreateMat (number_of_model_parameters, 1, CV_32FC1);
  CvMat * error                   = cvCreateMat( src_image->height, src_image->width, CV_32F);//error as a block
  
  const double src_norm = cvNorm(src_image_zeromean); 
  double dst_norm;

  // iteratively update map_matrix
  double rho      = 0;
  double correlation;
  double last_rho = - termination_eps;
  for (int i = 0; (i < number_of_iterations && fabs(rho - last_rho) >= termination_eps); i++)
  {  

    // warp-back portion of the dst_image to the coordinate space of the src_image
    if (warp_mode == WARP_MODE_AFFINE)
    {
      cvWarpAffine      (dst_image_float, dst_image_warped,      map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
      cvWarpAffine      (dst_gradient_x,  dst_gradient_x_warped, map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
      cvWarpAffine      (dst_gradient_y,  dst_gradient_y_warped, map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
    }
    else
    {
      cvWarpPerspective (dst_image_float, dst_image_warped,      map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
      cvWarpPerspective (dst_gradient_x,  dst_gradient_x_warped, map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
      cvWarpPerspective (dst_gradient_y,  dst_gradient_y_warped, map_matrix, CV_INTER_LINEAR+CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP);
    }
    
    // zero-mean shifting
    cvSubS (dst_image_warped, cvAvg (dst_image_warped), dst_image_warped);
    
    // calculate jacobian of warped-back destination image
    if (warp_mode == WARP_MODE_AFFINE)
		image_jacobian_affine(dst_gradient_x_warped, dst_gradient_y_warped, X, Y, jacobian);
    else
		image_jacobian_homo(dst_gradient_x_warped, dst_gradient_y_warped, X, Y, map_matrix, jacobian);

	// calculate Hessian and its inverse
	project_onto_jacobian(jacobian, jacobian, hessian);
    cvInvert (hessian, hessian_inv);
    
    // calculate enhanced correlation coefficiont (ECC) rho
    last_rho = rho;
	dst_norm = cvNorm(dst_image_warped);
	correlation = cvDotProduct (dst_image_warped,src_image_zeromean);
    rho = correlation/(src_norm*dst_norm);
//	printf("iteration contribution: %f\n", rho-last_rho);
//	printf("rho:%f\n", rho);

	project_onto_jacobian( jacobian, dst_image_warped, dst_image_projected);
	project_onto_jacobian( jacobian, src_image_zeromean, src_image_projected);

	
    cvGEMM (hessian_inv, dst_image_projected,     1, NULL, 0, dst_image_projected_hessian);

    // calculate the parameter lambda to account for illumination variation
    const double lambda_n = pow (dst_norm, 2.0) - cvDotProduct (dst_image_projected, dst_image_projected_hessian);
    const double lambda_d = correlation - cvDotProduct (src_image_projected, dst_image_projected_hessian);
    if (lambda_d <= 0.0)
    {
      // The execution was interrupted. The correlation value is going to be minimized.
      // Please check the warp initialization or the size of images.
      map_matrix = NULL;
      break;
    }
    const double lambda = (lambda_n/lambda_d);


    // estimate the update step delta_p
	cvScale (src_image_zeromean, error, lambda);
	cvSub  (error, dst_image_warped, error);
	//	cvAddWeighted(src_image_zeromean, lambda, dst_image_warped, -1.0, 0.0, error);
	
	project_onto_jacobian(jacobian, error, error_projected);
    cvGEMM (hessian_inv, error_projected, 1, NULL, 0, delta_p_flat);
    
    // update warping matrix
    update_warping_matrix (map_matrix, delta_p_flat);

  }

  // free intermediate buffers in reverse order
  cvReleaseMat (& error);
  cvReleaseMat (& delta_p_flat);
  cvReleaseMat (& error_projected);
  cvReleaseMat (& dst_image_projected_hessian);
  cvReleaseMat (& src_image_projected);
  cvReleaseMat (& dst_image_projected);
  cvReleaseMat (& hessian_inv);
  cvReleaseMat (& hessian);
  cvReleaseMat (& jacobian);
  cvReleaseMat (& dst_gradient_y_warped);
  cvReleaseMat (& dst_gradient_x_warped);
  cvReleaseMat (& dst_gradient_y);
  cvReleaseMat (& dst_gradient_x);
  cvReleaseMat (& dst_image_warped);
  cvReleaseMat (& dst_image_float);
  cvReleaseMat (& src_image_zeromean);
  
  // return updated map matrix
  return map_matrix;
}
